import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import NullLocator

palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)


def compute_color_for_labels(label):
    """
    Simple function that adds fixed color depending on the class
    """
    color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
    return tuple(color)



def draw_boxes(img, bbox, cls_name, confs, identities=None, offset=(0,0)):
    for i,box in enumerate(bbox):
        x1,y1,x2,y2 = [int(i) for i in box]
        x1 += offset[0]
        x2 += offset[0]
        y1 += offset[1]
        y2 += offset[1]
        # box text and bar
        id = int(identities[i]) if identities is not None else 0    
        color = compute_color_for_labels(id)
        # label = '{}{:d}'.format("", id)
        try:
            label = '{} {} {:.3f}'.format(id, cls_name, confs[i])
        except:
            label = '{} {} {:.3f}'.format(id, cls_name, confs[0])
        t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0]
        cv2.rectangle(img,(x1, y1),(x2,y2),color,3)
        cv2.rectangle(img,(x1, y1),(x1+t_size[0]+3,y1+t_size[1]+4), color,-1)
        cv2.putText(img,label,(x1,y1+t_size[1]+4), cv2.FONT_HERSHEY_PLAIN, 2, [255,255,255], 2)
    return img


# Bounding-box colors
cmap = plt.get_cmap("tab20b")
colors = [cmap(i) for i in np.linspace(0, 1, 20)]
def draw_boxes_with_plt(idx, save_path, img, bbox, cls_name, confs, identities=None, offset=(0,0)):
    plt.figure()
    fig, ax = plt.subplots(1)
    ax.imshow(img)
    for i,box in enumerate(bbox):
        x1,y1,x2,y2 = [int(i) for i in box]
        x1 += offset[0]
        x2 += offset[0]
        y1 += offset[1]
        y2 += offset[1]
        box_w = x2 - x1
        box_h = y2 - y1
        # box text and bar
        id = int(identities[i]) if identities is not None else 0
        color = compute_color_for_labels(id)
        color = tuple((np.array(color) /255.0))
        # label = '{}{:d}'.format("", id)
        try:
            label = '{} {} {:.3f}'.format(id, cls_name, confs[i])
        except:
            label = '{} {} {:.3f}'.format(id, cls_name, confs[0])
        # Create a Rectangle patch
        bbox = patches.Rectangle((x1, y1), box_w, box_h,
                                 linewidth=2, edgecolor=color, facecolor="none")
        # Add the bbox to the plot
        ax.add_patch(bbox)
        # Add label
        plt.text(
            x1,
            y1,
            s=label,
            color="white",
            verticalalignment="top",
            bbox={"color": color, "pad": 0},
        )

    # Save generated image with detections
    plt.axis("off")
    plt.gca().xaxis.set_major_locator(NullLocator())
    plt.gca().yaxis.set_major_locator(NullLocator())
    # filename = img_path.split("/")[-1].split(".")[0]
    plt.savefig("{}/{}.png".format(save_path, idx), bbox_inches="tight", pad_inches=0.0)
    plt.cla()
    plt.close('all')

    return "{}/{}.png".format(save_path, idx)


if __name__ == '__main__':
    for i in range(82):
        print(compute_color_for_labels(i))
